import re
import pickle
from tqdm import tqdm

def step2(file_path):
    result_sent, sent, record, idx = [], [], [], []
    count = 0
    with open(file_path,'r',encoding='gbk',errors='ignore') as file:
        content = file.readlines()
        for line in tqdm(content):
            ret = re.search('(\<\d+\>)(.+)',line.strip())
            if ret:
                flag = 0
                count = ret.group(1)
                sentence = ret.group(2)
                pos = sentence.split(' ')
                fenci = re.sub('\/[a-zA-Z]*','',sentence).split(' ')
                sent.append(''.join(i for i in fenci))
                num,i=[],0   
                for k in fenci:
                    num.append(list(range(i,i+len(k))))
                    i=i+len(k) 
            elif line.strip() == "====PGraph Beg====":
                continue
            elif line.strip() == "====PGraph End====":
                if flag == 0:
                    sent.append(num)
                    sent.append(idx)
                    sent.append(record)
                    sent.append(pos)
                    result_sent.append(sent)
                else:
                    flag = 0
                sent, record, idx = [], [], []
            elif len(line.strip()) != 0:
                head = line.strip().split(',')
                # if head[-1] in ['sbj_HX', 'obj_HX']:
                record.append((head[0],head[-1],head[2]))
                x = head[1].strip('()').split(':')
                y = head[3].strip('()').split(':')
                if len(set(x)) == 1 :
                    h = [int(i) for i in x[1:]]
                else:
                    m, n = int(x[0]), int(x[1])
                    h = list(range(m, m+(n-m)+1))
                if len(set(y)) == 1 :
                    t = [int(i) for i in y[1:]]
                else:
                    m, n = int(y[0]), int(y[1])
                    t = list(range(m, m+(n-m)+1))
                if (h in num) and (t in num):
                    idx.append((num.index(h),head[-1],num.index(t)))
                else:
                    flag = 1
    return result_sent

if __name__=="__main__":
    train = 'final/'+ 'train_data.pkl'
    val = 'final/'+ 'val_data.pkl'
    test = 'final/'+ 'test_data.pkl'
    # path = "../train_data/200w_all.out"
    path = "./train_data/new_50w.out"
    result_sent = step2(path)
    print(result_sent[:2])
    print(len(result_sent))
    # with open('final/200w_data.pkl','wb') as out_file:
    #     pickle.dump(result_sent,out_file) 
    # with open('final/200w_data.pkl','rb') as out_file:
    #     result_sent = pickle.load(out_file)  # len(result_sent) = 1999909
    with open(train,'wb') as out_file:
        pickle.dump(result_sent[:10000],out_file) 
    with open(val,'wb') as out_file:
        pickle.dump(result_sent[11000:12000],out_file) 
    with open(test,'wb') as out_file:
        pickle.dump(result_sent[12000:13000],out_file) 

    # with open('final/pred.pkl','wb') as out_file:
    #     pickle.dump(result_sent[1963000:1964000],out_file) 

    

